jax[cuda11_pip] @ https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
numpy
tqdm
scipy
gin-config
absl-py
flax
Pillow
optax
scikit-image
dm-pix
tensorflow
tensorflow-hub
mediapy